#include "stdafx.h"
#include "AudioCapThread.h"

#define REFTIMES_PER_SEC  10000000
#define REFTIMES_PER_MILLISEC  10000

#pragma comment(lib, "Avrt.lib")

CAudioCapThread::CAudioCapThread()
{
	m_dwDataSize = 0;
	m_bInit = false;
	//::InitializeCriticalSection(&m_cs);
}

CAudioCapThread::~CAudioCapThread()
{
}

void CAudioCapThread::SetDeiveType(int nType)
{
	m_nDeviceType = nType;
}


int CAudioCapThread::GetDeviceType()
{
	return m_nDeviceType;
}

void CAudioCapThread::OnCaptureData(LPBYTE pData, INT iDataLen)
{

	if (m_callback != NULL)
		m_callback(pData, iDataLen);
}


IMMDevice * CAudioCapThread::GetDefaultDevice(int nType)
{
	IMMDevice *pDevice = nullptr;

	IMMDeviceEnumerator *pMMDeviceEnumerator = nullptr;
	HRESULT hr = CoCreateInstance(
		__uuidof(MMDeviceEnumerator), nullptr, CLSCTX_ALL,
		__uuidof(IMMDeviceEnumerator),
		(void**)&pMMDeviceEnumerator);
	if (FAILED(hr))
		return nullptr;

	if (nType == CAudioCapThread::MICPHONE)
		hr = pMMDeviceEnumerator->GetDefaultAudioEndpoint((EDataFlow)eCapture, eConsole, &pDevice);
	else if (nType == CAudioCapThread::SPEAKER)
		hr = pMMDeviceEnumerator->GetDefaultAudioEndpoint((EDataFlow)eRender, eConsole, &pDevice);
	else
		pDevice = nullptr;

	if (pMMDeviceEnumerator)
	{
		pMMDeviceEnumerator->Release();
		pMMDeviceEnumerator = nullptr;
	}

	return pDevice;
}

bool CAudioCapThread::Init()
{
	if (m_bInit)
		return true;


	m_pDevice = GetDefaultDevice(m_nDeviceType);

	if (!m_pDevice)
		return false;

	m_dwDataSize = 0;

	m_bInit = true;

	return true;
}

DWORD CAudioCapThread::GetDataSize()
{
	DWORD dwSize = m_dwDataSize;
	return dwSize;
}

IMMDevice * CAudioCapThread::GetDevice()
{
	return m_pDevice;
}


void CAudioCapThread::Run()
{
	CoInitialize(NULL);

	HRESULT hr = 0;
	IAudioClient *pAudioClient = nullptr;
	WAVEFORMATEX *pWfx = nullptr;
	IAudioCaptureClient *pAudioCaptureClient = nullptr;
	DWORD nTaskIndex = 0;
	HANDLE hTask = nullptr;
	bool bStarted(false);
	int nDeviceType = 0;
	IMMDevice * pDevice = m_pDevice;

	REFERENCE_TIME hnsRequestedDuration = REFTIMES_PER_SEC;
	REFERENCE_TIME hnsActualDuration;

	UINT bufferFrameCount = 0;

	if (!pDevice )
		return ;

	do
	{
		hr = pDevice->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr, (void**)&pAudioClient);
		if (FAILED(hr))
			break;

		hr = pAudioClient->GetMixFormat(&pWfx);
		
		if (FAILED(hr))
			break;


		SaveFormat(pWfx);

		nDeviceType = GetDeviceType();

		if (nDeviceType == CAudioCapThread::MICPHONE)
			hr = pAudioClient->Initialize(AUDCLNT_SHAREMODE_SHARED, 0, 0, 0, pWfx, 0);
		else if (nDeviceType == CAudioCapThread::SPEAKER)
			hr = pAudioClient->Initialize(AUDCLNT_SHAREMODE_SHARED, AUDCLNT_STREAMFLAGS_LOOPBACK, 0, 0, pWfx, 0);
		else
			break;

		if (FAILED(hr))
			break;

		hr = pAudioClient->GetBufferSize(&bufferFrameCount);

		if (FAILED(hr))
			break;

		hr = pAudioClient->GetService(__uuidof(IAudioCaptureClient), (void**)&pAudioCaptureClient);

		if (FAILED(hr))
			break;

		hnsActualDuration = (double)REFTIMES_PER_SEC * bufferFrameCount / pWfx->nSamplesPerSec;

		if (nDeviceType == CAudioCapThread::MICPHONE)
			hTask = AvSetMmThreadCharacteristics(L"Audio", &nTaskIndex);
		else
			hTask = AvSetMmThreadCharacteristics(L"Capture", &nTaskIndex);

		if (!hTask)
			break;

		hr = pAudioClient->Start();
		if (FAILED(hr))
			break;

		bStarted = true;

		UINT32 uiNextPacketSize(0);
		BYTE *pData = nullptr;
		UINT32 uiNumFramesToRead;
		DWORD dwFlags;

		while (1)
		{
			if (IsStop())
				break;
			Sleep(hnsActualDuration / REFTIMES_PER_MILLISEC / 2);

			hr = pAudioCaptureClient->GetNextPacketSize(&uiNextPacketSize);
			if (FAILED(hr))
				break;

			while (uiNextPacketSize != 0)
			{
				hr = pAudioCaptureClient->GetBuffer(
					&pData,
					&uiNumFramesToRead,
					&dwFlags,
					nullptr,
					nullptr);

				if (FAILED(hr))
					break;

				if (dwFlags & AUDCLNT_BUFFERFLAGS_SILENT)
				{
					pData = NULL;
				}
		
				OnCaptureData(pData, uiNumFramesToRead * pWfx->nBlockAlign);	
	
				pAudioCaptureClient->ReleaseBuffer(uiNumFramesToRead);

				hr = pAudioCaptureClient->GetNextPacketSize(&uiNextPacketSize);

				if (FAILED(hr))
					break;
			}
		}

	} while (0);

	if (hTask)
	{
		AvRevertMmThreadCharacteristics(hTask);
		hTask = nullptr;
	}

	if (pAudioCaptureClient)
	{
		pAudioCaptureClient->Release();
		pAudioCaptureClient = nullptr;
	}

	if (pWfx)
	{
		CoTaskMemFree(pWfx);
		pWfx = nullptr;
	}

	if (pAudioClient)
	{
		if (bStarted)
		{
			pAudioClient->Stop();
		}

		pAudioClient->Release();
		pAudioClient = nullptr;
	}

}


void CAudioCapThread::SaveFormat(WAVEFORMATEX * wf)
{
	if (!wf)
		return;

	
	memcpy(&m_WaveFormat, wf, sizeof(WAVEFORMATEX));

	return;
}

WAVEFORMATEX * CAudioCapThread::GetWaveFormat()
{
	return &m_WaveFormat;
}



int CAudioCapThread::Start()
{
	if (!m_bInit)
		Init();
	TThreadRunable::Start();
	return 0;
}

void CAudioCapThread::Stop()
{
	if (m_bInit == false)
		return ;

	TThreadRunable::Stop();
	Join();
	if (m_pDevice)
	{
		m_pDevice->Release();
		m_pDevice = NULL;
	}
	m_bInit = false;

}